import java.util.Arrays;

public class Solution {
    public static void main(String[] args) {
        Solution s = new Solution();
        int[] vals = {1,3,2,1,3};
        int[][] edges = {{0,1},{0,2},{2,3},{2,4}};
        System.out.println(s.numberOfGoodPaths(vals, edges));
    }

    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        /**
         * 好路径的数目*/
        // 1 预处理
        int n = vals.length;
        initialize(n, vals);
        // -将边按照值大小排升序（本题核心逻辑，从小到大处理可以保证中间相连时不会出现大于的情况）
        Arrays.sort(edges, (a, b) -> Math.max(vals[a[0]], vals[a[1]]) - Math.max(vals[b[0]], vals[b[1]]));

        // 2 遍历使用并查集
        int ret = n;
        for(int[] e: edges) {
            // -两个集合最大值相等时才存在好路径
            int fa = find(e[0]);
            int fb = find(e[1]);
            if(maxValueOfTheSet[fa] == maxValueOfTheSet[fb]) {
                ret += theNumberOfMaxValueOfTheSet[fa] * theNumberOfMaxValueOfTheSet[fb];
            }
            // -合并两集合（需要同时更新最大值和最大值个数）
            union(fa, fb);
        }

        // 3 返回值
        return ret;
    }

    // 创建并查集
    public static int MAX_N;
    public static int[] father;
    public static int[] size;
    public static int[] stack;
    public static int[] maxValueOfTheSet;
    public static int[] theNumberOfMaxValueOfTheSet;

    public static void initialize(int n, int[] vals) {
        /**
         * 并查集初始化*/
        MAX_N = n;
        father = new int[MAX_N];
        size = new int[MAX_N];
        stack = new int[MAX_N];
        maxValueOfTheSet = new int[MAX_N];
        theNumberOfMaxValueOfTheSet = new int[MAX_N];
        for(int i = 0; i < n; i++) {
            father[i] = i;
            size[i] = 1;
            maxValueOfTheSet[i] = vals[i];
            theNumberOfMaxValueOfTheSet[i] = 1;
        }
    }

    public static int find(int a) {
        /**
         * 找到a元素所在集合的代表元素-扁平化*/
        int i = 0;
        while(a != father[a]) {
            stack[i++] = a;
            a = father[a];
        }
        while(i > 0) {
            father[stack[--i]] = a;
        }
        return a;
    }

    public static boolean isSameSet(int a, int b) {
        /**
         * 判断a、b元素是否处于同一集合中*/
        return find(a) == find(b);
    }

    public static void union(int a, int b) {
        /**
         * 合并a、b元素所在集合-小挂大*/
        int fa = find(a);
        int fb = find(b);
        if(fa != fb) {
            if(size[fa] >= size[fb]) {
                // -计算集合元素个数
                size[fa] += size[fb];
                // -更新最大值和最大值个数
                if(maxValueOfTheSet[fa] == maxValueOfTheSet[fb]) {
                    theNumberOfMaxValueOfTheSet[fa] += theNumberOfMaxValueOfTheSet[fb];
                } else {
                    if(maxValueOfTheSet[fa] < maxValueOfTheSet[fb]) {
                        maxValueOfTheSet[fa] = maxValueOfTheSet[fb];
                        theNumberOfMaxValueOfTheSet[fa] = theNumberOfMaxValueOfTheSet[fb];
                    }
                }
                // -小挂大
                father[fb] = fa;
            } else {
                size[fb] += size[fa];
                // -更新最大值和最大值个数
                if(maxValueOfTheSet[fa] == maxValueOfTheSet[fb]) {
                    theNumberOfMaxValueOfTheSet[fb] += theNumberOfMaxValueOfTheSet[fa];
                } else {
                    if(maxValueOfTheSet[fa] > maxValueOfTheSet[fb]) {
                        maxValueOfTheSet[fb] = maxValueOfTheSet[fa];
                        theNumberOfMaxValueOfTheSet[fb] = theNumberOfMaxValueOfTheSet[fa];
                    }
                }
                father[fa] = fb;
            }
        }
    }
}
